package client

import (
	"bufio"
	"crypto/rand"
	"crypto/sha1"
	"crypto/tls"
	"encoding/base64"
	"errors"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"
	"ws-demo-go/pkg/ws/common"
)

// Dialer 用于客户端拨号连接 WebSocket 服务端
type Dialer struct {
	HandshakeTimeout time.Duration // 握手超时时间（可选，默认10秒）
	TLSConfig        *tls.Config   // TLS 配置（wss:// 时使用）
}

// DefaultDialer 默认拨号器
var DefaultDialer = &Dialer{
	HandshakeTimeout: 10 * time.Second,
}

// Dial 拨号连接 WebSocket 服务端
// url: 服务端地址（如 "ws://localhost:8080/ws" 或 "wss://example.com/ws"）
// reqHeader: 额外的请求头（可选）
func (d *Dialer) Dial(wsURL string, reqHeader http.Header) (common.Conn, *http.Response, error) {
	// 解析 URL
	parsedURL, err := url.Parse(wsURL)
	if err != nil {
		return nil, nil, err
	}

	// 验证并转换协议（ws -> http，wss -> https）
	var isTLS bool
	switch parsedURL.Scheme {
	case "ws":
		parsedURL.Scheme = "http"
		isTLS = false
	case "wss":
		parsedURL.Scheme = "https"
		isTLS = true
	default:
		return nil, nil, common.ErrInvalidScheme
	}

	// 创建握手 HTTP 请求
	req, err := http.NewRequest(http.MethodGet, parsedURL.String(), nil)
	if err != nil {
		return nil, nil, err
	}

	// 设置 WebSocket 握手必需的请求头
	req.Header.Set("Upgrade", "websocket")
	req.Header.Set("Connection", "Upgrade")
	req.Header.Set("Sec-WebSocket-Key", generateWebSocketKey()) // 生成随机验证密钥
	req.Header.Set("Sec-WebSocket-Version", "13")               // 固定使用 RFC 6455 标准（版本13）

	// 添加用户自定义请求头（允许覆盖默认头，除了握手必需的）
	for k, v := range reqHeader {
		// 跳过握手必需的头（避免用户误改导致握手失败）
		switch strings.ToLower(k) {
		case "upgrade", "connection", "sec-websocket-key", "sec-websocket-version":
			continue
		}
		req.Header[k] = append(req.Header[k], v...)
	}

	// 建立底层连接（TCP 或 TLS）
	dialTimeout := d.HandshakeTimeout
	if dialTimeout <= 0 {
		dialTimeout = 10 * time.Second // 默认超时
	}
	netConn, err := d.dialUnderlying(parsedURL.Host, isTLS, dialTimeout)
	if err != nil {
		return nil, nil, err
	}

	// 执行握手并验证响应
	resp, err := d.doHandshake(netConn, req)
	if err != nil {
		netConn.Close() // 握手失败时关闭底层连接
		return nil, resp, err
	}

	// 握手成功，创建 WebSocket 连接
	wsConn := common.NewConn(netConn)
	return wsConn, resp, nil
}

// dialUnderlying 建立底层 TCP 或 TLS 连接
func (d *Dialer) dialUnderlying(addr string, isTLS bool, timeout time.Duration) (net.Conn, error) {
	if isTLS {
		// 建立 TLS 连接（wss://）
		tlsConfig := d.TLSConfig
		if tlsConfig == nil {
			tlsConfig = &tls.Config{
				InsecureSkipVerify: false, // 生产环境禁用跳过证书验证
			}
		}
		// 使用 tls.DialWithDialer 支持超时
		dialer := &net.Dialer{Timeout: timeout}
		return tls.DialWithDialer(dialer, "tcp", addr, tlsConfig)
	}

	// 建立 TCP 连接（ws://）
	return net.DialTimeout("tcp", addr, timeout)
}

// doHandshake 执行 WebSocket 握手并验证服务端响应
func (d *Dialer) doHandshake(netConn net.Conn, req *http.Request) (*http.Response, error) {
	// 设置握手超时（包含写请求和读响应的总时间）
	deadline := time.Now().Add(d.HandshakeTimeout)
	if err := netConn.SetDeadline(deadline); err != nil {
		return nil, err
	}
	defer netConn.SetDeadline(time.Time{}) // 清除超时，避免影响后续读写

	// 发送握手请求
	if err := req.Write(netConn); err != nil {
		return nil, errors.New("send handshake request failed: " + err.Error())
	}

	// 读取服务端响应（使用 bufio.Reader 符合 http.ReadResponse 要求）
	resp, err := http.ReadResponse(bufio.NewReader(netConn), req)
	if err != nil {
		return resp, errors.New("read handshake response failed: " + err.Error())
	}

	// 验证响应状态码（必须是 101 Switching Protocols）
	if resp.StatusCode != http.StatusSwitchingProtocols {
		return resp, errors.New("handshake failed: unexpected status code " + resp.Status)
	}

	// 验证 Upgrade 头（必须为 websocket）
	if !strings.EqualFold(resp.Header.Get("Upgrade"), "websocket") {
		return resp, errors.New("handshake failed: invalid Upgrade header")
	}

	// 验证 Connection 头（必须包含 Upgrade）
	if !strings.Contains(strings.ToLower(resp.Header.Get("Connection")), "upgrade") {
		return resp, errors.New("handshake failed: invalid Connection header")
	}

	// 验证 Sec-WebSocket-Accept 头（服务端需正确响应客户端的 Key）
	clientKey := req.Header.Get("Sec-WebSocket-Key")
	serverAccept := resp.Header.Get("Sec-WebSocket-Accept")
	if !validateWebSocketAccept(clientKey, serverAccept) {
		return resp, errors.New("handshake failed: invalid Sec-WebSocket-Accept")
	}

	return resp, nil
}

// generateWebSocketKey 生成符合协议的 Sec-WebSocket-Key（随机 16 字节，Base64 编码）
func generateWebSocketKey() string {
	var key [16]byte
	// 生成 16 字节随机数
	if _, err := rand.Read(key[:]); err != nil {
		panic("failed to generate WebSocket key: " + err.Error())
	}
	return base64.StdEncoding.EncodeToString(key[:])
}

// validateWebSocketAccept 验证服务端返回的 Sec-WebSocket-Accept 是否正确
// 算法：将客户端 Key + WebsocketGUID 拼接后做 SHA-1 哈希，再 base64 编码
func validateWebSocketAccept(clientKey, serverAccept string) bool {
	hash := sha1.Sum([]byte(clientKey + common.WebsocketGUID))
	expectedAccept := base64.StdEncoding.EncodeToString(hash[:])
	return serverAccept == expectedAccept
}
